using LinearAlgebra, SpecialFunctions

function LegendreGaussLobatto(p)
    """
    Compute the p'th order LGL solution points at reference domain: [-1,1]
    """
    x = zeros(Float64, p+1)

    # p=1, x = [-1, 1]
    if p == 1
        x[1] = -1.0
        x[2] =  1.0
        return x
    end

    # p = 2, x = [-1, 0, 1]
    if p == 2
        x[1]= -1.0
        x[2]=  0.0
        x[3]=  1.0
        return x
    end

    h = @. 2*(0:p-2) + 2
    J = zeros(p-1, p-1)
    @inbounds A = @. sqrt((1:p-2)*((1:p-2)+2)*((1:p-2)+1).*((1:p-2)+1)/(h[1:p-2]+1)/(h[1:p-2]+3))
    @inbounds J[1:p-2,2:p-1] = Diagonal(2.0 ./ (h[1:p-2] .+ 2).*A)
    @inbounds J[1,1] = 0.0
    @inbounds J = J + J'

    @inbounds D, V = eigen(J)
    @inbounds i = sortperm(D)
    x = [-1.0; D[i]; 1.0]
end

function LegendreP(p, x)
    """
    Evaluate p'th order orhonormal Legendre Polynomial at x
    """
    # Initial values for P(x) at x
    P = zeros(p+1, length(x))
    @. P[1,:] = sqrt(1.0/2.0)

    # p = 1
    if  p == 0
        return P'
    end

    # p = 2
    @inbounds P[2,:] = sqrt(3.0/2.0)*x
    if  p == 1
        return P[p+1,:]
    end

    # Forward recurrence using the symmetry of the recurrence for p>=2
    a₀ = sqrt(1.0/3.0)
    @simd for i=1:p-1
        a₁ = 2/(2*i+2)*sqrt((i+1)*(i+1)*(i+1)*(i+1)/(2*i+1)/(2*i+3))
        @inbounds P[i+2,:] = 1/a₁*(-a₀*P[i,:] + x.*P[i+1,:])
        a₀ = a₁
    end

    return P[p+1,:]
end

function GradLegendreP(p, x)
    """
    Evaluate the derivative of the p'th order Legendre polynomial at x
    """
    P₀ = zeros(length(x),1)
    if p > 0
        @inbounds P₁ = -p*x.*LegendreP(p, x) + p*sqrt((2*p+1)/(2*p-1))*LegendreP(p-1,x)
        @inbounds P₃ = x.^(p+1)*p*(p+1)/2*sqrt((2*p+1)/2)
        @inbounds I  = abs.(abs.(x) .- 1) .> 10*2.2204e-16
        @inbounds x₁ = x .* I
        @inbounds P₀ = broadcast(~,I).*P₃ + I.*P₁./(1.0 .- x₁.^2)
    end
    return P₀
end

function JacobiP(p, α, β, x)
    """
    Evaluate Jacobi Polynomial of type (α, β) > -1 (α+ β <> -1) at points x for order p and returns P[1:length( x))]
    Note: They are normalized to be orthonormal.
    """
    P = zeros( p+1,length(x))

    # Initial values for P(x) at x
    γ₀ = 2^(α+β+1)/(α+β+1)*gamma(α+1)*gamma(β+1)/gamma(α+β+1)
    @. P[1,:] = 1/sqrt(γ₀)

    # p = 1
    if  p == 0
        return P'
    end

    # p = 2
    γ₁ = (α+1)*(β+1)/(α+β+3)* γ₀
    @inbounds P[2,:] = ((α+β+2)* x/2 .+ (α-β)/2)/sqrt(γ₁)
    if  p == 1
        return P[p+1,:]
    end

    # Repeat value in recurrence
    a₀ = 2/(2+α+ β)*sqrt((α+1)*( β+1)/(α+ β+3))
    # Forward recurrence using the symmetry of the recurrence.
    for i=1:p-1
        c = 2*i+α+β
        a₁ = 2/(c+2)*sqrt((i+1)*(i+1+α+β)*(i+1+α)*(i+1+β)/(c+1)/(c+3))
        b₀ = -(α^2-β^2)/c/(c+2)
        @inbounds P[i+2,:] = 1/a₁*( -a₀* P[i,:] .+ ( x .- b₀).* P[i+1,:])
        a₀ =a₁
    end
    return P[p+1,:]
end

function GradJacobiP(p, α, β, x)
    """
    Evaluate the derivative of the Jacobi polynomial of type (α, β)>-1,
    at x for order p and returns P[1:length(r))]
    """
    P = zeros(length(x),1)

    # p > 1
    if p > 0
        P = sqrt(p*(p+α+β+1))*JacobiP(p-1,α+1,β+1,x)
    end
    P
end

function Warpfactor(p, x)
    """
    Compute scaled warp function at order p based on x interpolation nodes
    """
    x₀ = LegendreGaussLobatto(p)           # LGL points at [-1,1]
    x₁ = range(-1,stop = 1,length=p+1)     # equit-distance points at [-1,1]

    # Compute Vandermonde matrix V based on
    V = Vandermonde1D(p, x₁)

    # Evaluate Lagrange polynomial at x
    P = zeros(p+1,length(x))
    for i = 1:p+1
        P[i,:] = LegendreP(i-1,x)'
    end

    # Compute w factor
    w = (V'\P)'*(x₀ - x₁)

    # Scale factor
    s = 1.0 .- ((abs.(x) .< 1.0 .- 1.0e-12).*x).^2
    w = w./s + w.*((abs.(x) .< 1.0 .- 1.0e-12) .- 1.0)
end

function Nodes(p)
    """
    Purpose: Compute (x,y) nodes in equilateral triangle for polynomial of order p
    """
    αₒ = [0.0000; 0.0000; 1.4152;
          0.1001; 0.2751; 0.9800;
          1.0999; 1.2832; 1.3648;
          1.4773; 1.4959; 1.5743;
          1.5770; 1.6223; 1.6258]

    # Set optimized parameter α, depending on order p
    if p < 16
        α = αₒ[p]
    else
        α = 5.0/3.0
    end

    # Total number of Nodes
    Np = convert(Int,(p+1)*(p+2)/2)

    # Create equidistributed nodes on equilateral triangle
    x₁ = zeros(Np,1)
    x₂ = zeros(Np,1)
    x₃ = zeros(Np,1)

    i₀ = 1
    for i = 1:p+1
        for j = 1:p+2-i
            x₁[i₀] = (i-1)/p
            x₃[i₀] = (j-1)/p
            x₂[i₀] = 1.0.-(x₁[i₀] + x₃[i₀])
            i₀ = i₀ + 1
        end
    end

    x = x₃ - x₂
    y = (2*x₁ - x₂ - x₃)/sqrt(3.0)

    # Compute blending function at each node for each edge
    b₁ = 4*x₂.*x₃
    b₂ = 4*x₁.*x₃
    b₃ = 4*x₁.*x₂

    # Amount of warp for each node, for each edge
    w₁ = Warpfactor(p,x₃-x₂)
    w₂ = Warpfactor(p,x₁-x₃)
    w₃ = Warpfactor(p,x₂-x₁)

    # Combine blend & warp
    w₁ = b₁.*w₁.*(1 .+ (α*x₁).^2)
    w₂ = b₂.*w₂.*(1 .+ (α*x₂).^2)
    w₃ = b₃.*w₃.*(1 .+ (α*x₃).^2)

    # Accumulate deformations associated with each edge
    x = x + 1*w₁ + cos(2*pi/3)*w₂ + cos(4*pi/3)*w₃
    y = y + 0*w₁ + sin(2*pi/3)*w₂ + sin(4*pi/3)*w₃

    x, y
end

function xytors(x,y)
    """
    From (x,y) in equilateral triangle to (r,s) coordinates in standard triangle.
    """
    α = (sqrt(3.0)*y .+ 1.0)/3.0
    β = (-3.0*x .- sqrt(3.0)*y .+ 2.0)/6.0
    γ = ( 3.0*x .- sqrt(3.0)*y .+ 2.0)/6.0

    r = -α-β+γ
    s =  α-β-γ
    r,s
end

function rstoab(r,s)
    """
    From (r, s) to (a, b)
    """
    Np = length(r)
    a = zeros(Np,1)

    for n = 1:Np
        if s[n] != 1
            a[n] = 2*(1 .+ r[n]) /(1 .- s[n]) .- 1
        else
            a[n] = -1
        end
    end
    b = s

    a, b
end

function Simplex2DP(a,b,i,j)
    """
    Purpose: Evaluate 2D orthonormal polynomial on simplex at (a,b) of order (i,j).
    """
    P₁ = JacobiP(i,0,0,a)
    P₂ = JacobiP(j,2*i+1,0,b)
    P  = sqrt(2.0)*P₁.*P₂.*(1 .- b).^i
end

function GradSimplex2DP(a,b,i,j)
    """
    Return the derivatives of the modal basis (i,j) on the 2D simplex at (a,b).
    """
    fa = JacobiP(i,0,0,a)
    dfa = GradJacobiP(i,0,0,a)
    gb = JacobiP(j,2*i+1,0,b)
    dgb = GradJacobiP(j,2*i+1,0,b)

    dmodedr = dfa.*gb
    if i > 0
        dmodedr = dmodedr.*((0.5*(1 .- b)).^(i - 1))
    end

    dmodeds = dfa.*(gb.*(0.5*(1 .+ a)));
    if i > 0
        dmodeds = dmodeds.*((0.5*(1 .- b)).^(i - 1));
    end

    tmp = dgb.*((0.5*(1 .- b)).^i)
    if i > 0
        tmp = tmp - 0.5*i*gb.*((0.5*(1 .- b)).^(i - 1))
    end
    dmodeds = dmodeds+fa.*tmp

    # Normalize
    dmodedr = 2^(i + 0.5)*dmodedr; dmodeds = 2^(i + 0.5)*dmodeds

    dmodedr, dmodeds
end

function Vandermonde1D(p, x)
    """
    Initialize the 1D Vandermonde Matrix, V_{ij} = phi_j(r_i)
    """
    V = zeros(length(x),p+1)
    @simd for i in 1:p+1
            V[:,i] = LegendreP(i-1,x)
    end
    V
end

function Vandermonde2D(p, r, s)
    """
    Initialize the 2D Vandermonde Matrix,  V_{ij} = phi_j(r_i, s_i);
    """
    Np = convert(Int, (p+1)*(p+2)/2)
    V = zeros(length(r),Np)

    # Transfer to (a,b) coordinates
    a, b = rstoab(r,s)

    # Build the Vandermonde Matrix
    i₀ = 1
    for i = 0:p
        for j = 0:p-i
            V[:,i₀] = Simplex2DP(a,b,i,j)
            i₀ = i₀ + 1
        end
    end
    V
end

function GradVandermonde2D(p, r, s)
    """
    Initialize the gradient of the modal basis (i,j) at (r,s) at order N
    """
    Np = convert(Int,(p+1)*(p+2)/2)
    Vr = zeros(length(r),Np)
    Vs = zeros(length(r),Np)

    # find tensor-product coordinates
    a, b = rstoab(r,s)

    # Initialize Matrix
    i₀ = 1
    for i = 0:p
        for j = 0:p-i
            Vr[:,i₀], Vs[:,i₀] = GradSimplex2DP(a,b,i,j)
            i₀ = i₀ + 1
        end
    end
    Vr, Vs
end

function DiffMatrix1D(p, x, V)
    """
    Initialize the differentiation matrices D
    """
    Vx = zeros(length(x),p+1)
    for i=0:p
       Vx[:,i+1] = GradLegendreP(i, x)
    end

    D = Vx/V
    return D
end

function DiffMatrix2D(p, r, s, V)
    """
    Initialize the gradient of the modal basis (i,j) at (r,s) at order N
    """
    Vr, Vs = GradVandermonde2D(p,r,s)
    Dr = Vr/V
    Ds = Vs/V
    Dr, Ds
end

mutable struct Operator
    """
    A object contain element-wise operators.
    """
    P::Int64                # order of interpolation polynomials
    Np::Int64               # number of points at each element
    x₀::Array{Float64,1}    # solution points at [-1,1]
    r::Array{Float64,2}
    s::Array{Float64,2}

    V::Array{Float64,2}
    Dr::Array{Float64,2}
    Ds::Array{Float64,2}
    Drw::Array{Float64,2}
    Dsw::Array{Float64,2}

    function Operator(p)
        self = new()

        self.P = p
        self.Np = (p+1)*(p+2)/2
        self.x₀ = LegendreGaussLobatto(p)
        x, y = Nodes(p)
        self.r, self.s = xytors(x,y)
        V = Vandermonde2D(p,self.r,self.s); self.V = V
        self.Dr, self.Ds = DiffMatrix2D(p,self.r,self.s,self.V)

        # Compute weak operators (could be done in preprocessing to save time)
        Vr, Vs = GradVandermonde2D(p, self.r, self.s)
        @inbounds self.Drw = (V*Vr')/(V*V')
        @inbounds self.Dsw = (V*Vs')/(V*V')
        return self
    end
end
